Explaining a nlp model’s layer attentions

See BertViz[Vig19]

from explainer.explainers import layer_attentions
from transformers import BertTokenizer, BertModel, utils
utils.logging.set_verbosity_error()  # Suppress standard warnings

model_version = 'bert-base-uncased'
model = BertModel.from_pretrained(model_version, output_attentions=True)
tokenizer = BertTokenizer.from_pretrained(model_version)
sentence_a = "The cat sat on the mat"
sentence_b = "The cat lay on the rug"
inputs = tokenizer.encode_plus(sentence_a, sentence_b, return_tensors='pt')
input_ids = inputs['input_ids']
token_type_ids = inputs['token_type_ids']
attention = model(input_ids, token_type_ids=token_type_ids)[-1]
sentence_b_start = token_type_ids[0].tolist().index(1)
input_id_list = input_ids[0].tolist() # Batch index 0
tokens = tokenizer.convert_ids_to_tokens(input_id_list) 

import entry_point
explainer = entry_point.entry_point(attention, tokens, sentence_b_start)
explainer()

Layer: Attention: All Sentence A -> Sentence A Sentence B -> Sentence B Sentence A -> Sentence B Sentence B -> Sentence A